import json
import os
import time
from eth_account import Account
from eth_typing import ChecksumAddress
from web3 import HTTPProvider, Web3
import logging

PRODUCTION_RPC_URL = "http://18.143.239.137:8545"

TOKEN_TO_ADDRESS = {
    "ETH": "0x0000000000000000000000000000000000000000",
    "WETH": "0xC02aaA39b223FE8D0A0e5C4F27eAD9083C756Cc2",
    "DAI": "0x6B175474E89094C44Da98b954EedeAC495271d0F",
    "USDC": "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48",
    "BAT": "0x0D8775F648430679A709E98d2b0Cb6250d2887EF",
    "WBTC": "0x2260FAC5E5542a773Aa44fBCfeDf7C193bc2C599",
    "UNI": "0x1f9840a85d5aF5bf1D1762F925BDADdC4201F984",
    "SHIBA2": "0xA0335820dC549dBfae5b8D691331CadfCA7026E0",
    "USDT": "0xdAC17F958D2ee523a2206206994597C13D831ec7",
    "CCX": "0x1D6405138A335cE5FD7364086334EfB3E4F28B59",
}

PRIVATE_KEY = None

deployment_address = {
  'UniswapV2Factory': '0x5C69bEe701ef814a2B6a3EDD4B1652CB9cc5aA6f',
  'UniswapV2Router02': '0x7a250d5630B4cF539739dF2C5dAcb4c659F2488D',
}


def _load_abi(name: str) -> str:
    path = "/home/yeming/workspace/coin/python/coin/exchange/uniswap_v2/abi"
    with open(os.path.join(path, f"{name}.abi")) as f:
        abi: str = json.load(f)
    return abi


class UniswapV2NativePrivateClient(object):
    def __init__(self):
        self.account: Account = Account.from_key(PRIVATE_KEY)
        self.wallet_address: ChecksumAddress = self.account.address
        self.uniswap_instance = Web3(HTTPProvider(
            PRODUCTION_RPC_URL
        ))
        self.currencies = ["ETH", "WETH", "CCX", "USDT"]
        self.token_list = [TOKEN_TO_ADDRESS[symb] for symb in self.currencies]
        self.current_nonce = self.uniswap_instance.eth.get_transaction_count(
            self.wallet_address
        )
        print(self.wallet_address)

    def _router(self):
        return self.uniswap_instance.eth.contract(
            address=Web3.to_checksum_address(deployment_address["UniswapV2Router02"]),
            abi=_load_abi("router02")
        )
    router = property(fget=_router, fset=None, fdel=None, doc=None)

    def swap_out(
        self,
        path,
        amount_out,
        amount_in_maximum=((1 << 256) - 1),
        deadline=None,
    ):
        if deadline is None:
            deadline = int(time.time()) + 2 * 60

        assert len(path) >= 2
        args = None
        tx_args = {
            "nonce": self.current_nonce,
            "from": self.wallet_address,
            "type": "0x2",
            "maxFeePerGas": self.uniswap_instance.eth.max_priority_fee
            + (2 * self.uniswap_instance.eth.get_block("latest")["baseFeePerGas"]),
            "maxPriorityFeePerGas": self.uniswap_instance.eth.max_priority_fee,
        }

        exec_func = self.router.functions.swapTokensForExactTokens

        if path[-1] == TOKEN_TO_ADDRESS["ETH"]:
            path[-1] = TOKEN_TO_ADDRESS["WETH"]
            exec_func = self.router.functions.swapTokensForExactETH

        if path[0] == TOKEN_TO_ADDRESS["ETH"]:
            path[0] = TOKEN_TO_ADDRESS["WETH"]
            ETH_USED = amount_in_maximum
            exec_func = self.router.functions.swapETHForExactTokens
            tx_args |= {
                "value": ETH_USED
            }
            args = (amount_out, path, self.wallet_address, deadline)

        if args is None:
            args = (amount_out, amount_in_maximum, path, self.wallet_address, deadline)

        transaction = exec_func(
            *args
        ).build_transaction(tx_args)
        self.current_nonce += 1
        transaction["gas"] = int(
            self.uniswap_instance.eth.estimate_gas(transaction) * 1.2
        )
        signed_txn = self.uniswap_instance.eth.account.sign_transaction(
            transaction, private_key=self.account.key
        )
        transaction = self.uniswap_instance.eth.send_raw_transaction(
            signed_txn.rawTransaction
        )
        receipt = self.uniswap_instance.eth.wait_for_transaction_receipt(transaction)
        return receipt


    def swap_in(
        self,
        path,
        amount_in,
        amount_out_minimum=0,
        deadline=None,
    ):
        if deadline is None:
            deadline = int(time.time()) + 2 * 60

        assert len(path) >= 2
        args = None
        tx_args = {
            "nonce": self.current_nonce,
            "from": self.wallet_address,
            "type": "0x2",
            "maxFeePerGas": self.uniswap_instance.eth.max_priority_fee
            + (2 * self.uniswap_instance.eth.get_block("latest")["baseFeePerGas"]),
            "maxPriorityFeePerGas": self.uniswap_instance.eth.max_priority_fee,
        }

        exec_func = self.router.functions.swapExactTokensForTokens
        if path[-1] == TOKEN_TO_ADDRESS["ETH"]:
            path[-1] = TOKEN_TO_ADDRESS["WETH"]
            exec_func = self.router.functions.swapExactTokensForETH
        
        if path[0] == TOKEN_TO_ADDRESS["ETH"]:
            path[0] = TOKEN_TO_ADDRESS["WETH"]
            ETH_USED = amount_in
            exec_func = self.router.functions.swapExactETHForTokens
            tx_args |= {
                "value" : ETH_USED
            }
            args = (amount_in, path, self.wallet_address, deadline)

        if args is None:
            args = (amount_in, amount_out_minimum, path, self.wallet_address, deadline)

        print(args)
        print(tx_args)
        transaction = exec_func(
            *args            
        ).build_transaction(
            tx_args
        )
        self.current_nonce += 1
        transaction["gas"] = int(
            self.uniswap_instance.eth.estimate_gas(transaction) * 1.2
        )
        signed_txn = self.uniswap_instance.eth.account.sign_transaction(
            transaction, private_key=self.account.key
        )
        transaction = self.uniswap_instance.eth.send_raw_transaction(
            signed_txn.rawTransaction
        )
        receipt = self.uniswap_instance.eth.wait_for_transaction_receipt(transaction)
        return receipt


    def approve_contracts(self):
        # contracts = self.get_token_contracts() # USDT and CCX
        usdc_contract = self.uniswap_instance.eth.contract(
            address=Web3.to_checksum_address(TOKEN_TO_ADDRESS["USDC"]),
            abi=_load_abi("erc20")
        )
        
        ccx_contract = self.uniswap_instance.eth.contract(
            address=Web3.to_checksum_address(TOKEN_TO_ADDRESS["CCX"]),
            abi=_load_abi("erc20")
        )
        
        contracts = [usdc_contract, ccx_contract]
        
        for instance in contracts:
            tx = instance.functions.approve(
                deployment_address["UniswapV2Router02"], 1000000
            ).build_transaction(
                {
                    "nonce": self.current_nonce,
                    "from": self.wallet_address,
                    "type": "0x2",
                    "maxFeePerGas": self.uniswap_instance.eth.max_priority_fee
                    + (
                        2
                        * self.uniswap_instance.eth.get_block("latest")["baseFeePerGas"]
                    ),
                    "maxPriorityFeePerGas": self.uniswap_instance.eth.max_priority_fee,
                }
            )
            self.current_nonce += 1
            tx["gas"] = int(self.uniswap_instance.eth.estimate_gas(tx) * 1.2)
            logging.info(tx)
            signed_txn = self.uniswap_instance.eth.account.sign_transaction(
                tx, private_key=self.account.key
            )
            transaction = self.uniswap_instance.eth.send_raw_transaction(
                signed_txn.rawTransaction
            )
            recipt = self.uniswap_instance.eth.wait_for_transaction_receipt(transaction)
            logging.info(recipt)


def test():
    account = Account.from_key(PRIVATE_KEY)
    print(account.address)
    obj: UniswapV2NativePrivateClient = UniswapV2NativePrivateClient()
    obj.approve_contracts()
    obj.swap_in(
        [
            TOKEN_TO_ADDRESS["USDC"],
            TOKEN_TO_ADDRESS["WETH"],
            TOKEN_TO_ADDRESS["CCX"],
        ],
        100000
    )
    
if '__main__' == __name__:
    test()
